import pandas as pd
import numpy as np
import plotly.express as px


def merge_rows_by_dict(df, merge_dict):
    """
    根据字典规则将 DataFrame 的若干行合并为新行。

    参数:
    df (pd.DataFrame): 输入的 DataFrame。
    merge_dict (dict): 合并规则字典，键为要合并的行名列表，值为合并后的新行名。

    返回:
    pd.DataFrame: 合并后的 DataFrame。
    """
    tree_dict = {}
    for new_row_name, old_row_names in merge_dict.items():
        li = tree_dict.setdefault(old_row_names, [])
        li.append(new_row_name)
    container = {}
    for name, li in tree_dict.items():
        se = df.loc[li, :].sum()
        container[name] = se
    merge_df = pd.DataFrame(container)

    return merge_df


def transform(df, merge_dict, by_row=True, by_col=True):
    # 合并行，不合并列
    if by_row and not by_col:
        # 此时，行合并到新的行，同时行变成列,需要转置
        return merge_rows_by_dict(df, merge_dict).T
    # 合并列，合并行
    if by_col and by_row:
        df = merge_rows_by_dict(df, merge_dict)
        return merge_rows_by_dict(df, merge_dict)

    if by_col and not by_row:
        return merge_rows_by_dict(df.T, merge_dict)

    if not by_row and not by_col:
        return df


def calculate_leontief_inverse(A):
    """
    计算里昂惕夫逆矩阵。

    参数:
    A (pd.DataFrame): 直接消耗系数矩阵，类型为 Pandas 的 DataFrame。

    返回:
    pd.DataFrame: 里昂惕夫逆矩阵，类型为 Pandas 的 DataFrame，索引和列名与输入矩阵 A 相同。
    """
    # 创建与 A 相同形状的单位矩阵
    I = pd.DataFrame(np.eye(len(A)), index=A.index, columns=A.columns)
    # 计算 I - A
    I_minus_A = I - A
    # 计算 (I - A) 的逆矩阵
    leontief_inverse = pd.DataFrame(np.linalg.inv(I_minus_A.values), index=A.index, columns=A.columns)
    return leontief_inverse


class InputOutputTable:
    def __init__(self, year, country, df_collections, target_category , unit='亿元'):
        self.year = year
        self.country = country
        self.df_collections = df_collections
        self.target_category = target_category
        self.unit = unit
        self.merge_dict = self.parse_merge_dict()
        self.q1, self.q2, self.q3, self.A, self.L, self.B, self.F, self.E = self.ini_table()
        self.assert_q2()
        self.set_quick_attr() # 设置一些快速获取的属性
        self.D = self.income_coef()
        self.APL = self.calculate_apl()


    def parse_merge_dict(self):
        df = self.df_collections['merge']
        df = df.dropna(axis=0, how='any')
        merge_dict = dict(zip(df['source'], df[self.target_category]))
        return merge_dict

    def calculate_apl(self, max_iter=100):
        """
        计算投入产出表的 APL（Average Propagation Length）系数。

        参数:
        A (pd.DataFrame): 直接消耗系数矩阵。
        L (pd.DataFrame): 里昂惕夫逆矩阵。
        max_iter (int): 最大迭代次数，用于近似计算无穷级数，默认为 100。

        返回:
        pd.DataFrame: APL 系数矩阵，索引和列名与输入矩阵 A 相同。
        """
        # 初始化求和结果矩阵
        A = self.A
        L = self.L
        sum_matrix = pd.DataFrame(0, index=A.index, columns=A.columns)
        A_k = A.copy()
        for k in range(1, max_iter + 1):
            # 计算 k * (A^k)
            k_A_k = k * A_k
            # 累加结果
            sum_matrix += k_A_k
            # 计算下一次的 A^k
            A_k = A_k.dot(A)
        # 计算 APL 系数
        apl_matrix = sum_matrix / L
        return apl_matrix

    def ini_table(self):
        q1 = self.df_collections['Q1']
        q2 = self.df_collections['Q2']
        q3 = self.df_collections['Q3']
        if self.unit == '亿元':
            q1 = q1 / 10000
            q2 = q2 / 10000
            q3 = q3 / 10000
        # 中间投入矩阵，行列都要合并
        q1 = transform(q1, self.merge_dict, by_row=True, by_col=True)
        # 支出法核算矩阵，仅行合并
        q2 = transform(q2, self.merge_dict, by_row=True, by_col=False)
        # 生产法矩阵，列合并
        q3 = transform(q3, self.merge_dict, by_row=False, by_col=True)

        # 计算直接消耗性系数
        total_output = q2['总产出']
        A = q1 / total_output
        # 绝对数据用小数点两位，相对数据用小数点6位
        # 计算完全消耗系数
        # - 计算列昂惕夫逆矩阵
        L = calculate_leontief_inverse(A)
        # 计算单位阵
        I = pd.DataFrame(np.eye(len(A)), index=A.index, columns=A.columns)
        B = L - I

        # 计算影响力系数
        average = L.sum().sum() / L.shape[0]
        F = L.sum(axis=0) / average

        # 计算感应度系数
        E = L.sum(axis=1) / average

        return q1.round(2), q2.round(2), q3.round(2), A.round(6), L.round(6), B.round(6), F.round(6), E.round(6)

    def assert_q2(self):
        normal_q2_cols = [
            '中间使用合计','农村居民消费支出','城镇居民消费支出','居民消费支出','政府消费支出','消费支出合计','固定资本形成总额','存货变动','资本形成总额','出口','最终使用合计','进口','总产出'
            ]
        if '误差' in self.q2.columns:
            print(self.year)
            col_map = {
                '农村居民消费':'农村居民消费支出',
                '城镇居民消费':'城镇居民消费支出',
                '政府消费':'政府消费支出',
                '消费小计':'消费支出合计',
                '存货增加':'存货变动'}
            new_cols = [col_map[s] if s in col_map else s for s in self.q2.columns]
            self.q2.columns = new_cols
            self.q2['居民消费支出'] = self.q2['农村居民消费支出'] + self.q2['城镇居民消费支出']
            del self.q2['误差']
        
            print(f"year {self.year} q2 cols reset", normal_q2_cols)

            self.q2 = self.q2.loc[:, normal_q2_cols]

    def set_quick_attr(self):
        q3_idx_code = ['inter_input', 'wage', 'tax', 'asset_dep', 'profit', 'value_added', 'total_input']
        q3_maps = dict(zip(q3_idx_code, self.q3.index.to_list()))
        for key, value in q3_maps.items():
            setattr(self, key, self.q3.loc[value, :])
        q2_col_code = ['inter_use', 'rural_con', 'urban_con', 'household_con', 'gov_con', 'total_con', 
                       'gfcf', 'change_inv', 'gcf', 'exp', 'final_use', 'imp', 'total_output']
        q2_maps = dict(zip(q2_col_code, self.q2.columns.to_list()))
        for key, value in q2_maps.items():
            setattr(self, key, self.q2[value])
        self.q2_map = q2_maps
        self.q3_map = q3_maps

    def PIC(self, demand_matrix):
        # p = L * D
        """
        生产诱发度，Production Inducement Coefficient 是衡量某一最终需求项目，如消费投资，出口对各产业部门生产活动带动程度的指标。它反映了当某种最终需求增加一个单位时，会诱发各个产业部门产出增加的数量。
        """
        P = self.L.dot(demand_matrix)
        return P

    def PIC_q2(self, name):
        """
        计算最终需求拉动的生产诱发度。

        参数:
        name (str): 最终需求项目的名称，必须是 q2 数据框的列名。

        返回:
        pd.Series: 包含每个产业部门的生产诱发度值的 Series，索引为产业部门名称。
        """

        vector = self.q2[name]
        vector_ratio = vector / vector.sum()
        vector_df = pd.DataFrame(vector_ratio)
        return self.PIC(vector_df)

    def income_coef(self, production_sub_items='劳动者报酬'):
        """
        计算收入分配系数矩阵。
        收入分配系数矩阵描述了各产业部门的劳动者报酬在总投入中的占比。

        返回:
        pd.DataFrame: 收入分配系数矩阵，对角线元素为各产业部门劳动者报酬占总投入的比例，索引和列名与直接消耗系数矩阵 self.A 相同。
        """
        # 创建一个与直接消耗系数矩阵 self.A 相同形状的单位矩阵
        # 该矩阵将作为基础，后续更新对角线元素为收入分配系数
        D = pd.DataFrame(np.eye(len(self.A)), index=self.A.index, columns=self.A.columns)
        # 遍历直接消耗系数矩阵的每一行（即每个产业部门）
        for i in range(len(self.A)):
            # 从生产法矩阵 self.q3 中获取第 i 列的数据
            # 该列数据包含了对应产业部门的各项投入信息
            se = self.q3.iloc[:, i]
            # 计算该产业部门劳动者报酬在总投入中的比例
            ratio = se[production_sub_items] / se['总投入']
            # 将计算得到的比例更新到单位矩阵 D 的对应对角线位置
            D.iloc[i, i] = ratio
        return D

    def IIC_q2(self, name, production_sub_items='劳动者报酬'):

        """
        计算收入诱导度
        收入诱导度 = 收入分配系数矩阵 × 列昂惕夫逆矩阵 × 最终需求向量

        参数:
        name (str): 最终需求项目的名称，必须是 q2 数据框的列名。

        返回:
        pd.Series: 包含每个产业部门的收入诱导系数值的 Series，索引为产业部门名称。
        """
        # D为收入分配系数矩阵
        # L为列昂惕夫逆矩阵
        # vector_df为最终需求向量
        D = self.income_coef(production_sub_items)
        L = self.L
        if name not in self.q2.columns:
            print(f"year {self.year} 投入产出表，最终需求列向量包括:", self.q2.columns)
        vector = self.q2[name]
        vector_ratio = vector / vector.sum() # 形成最终需求向量，总和为1，占比与原始数据一致
        vector_df = pd.DataFrame(vector_ratio)
        return D.dot(L).dot(vector_df)

    def IIC_sector(self, sector, production_sub_items='劳动者报酬'):

        """
        计算收入诱导度
        收入诱导度 = 收入分配系数矩阵 × 列昂惕夫逆矩阵 × 最终需求向量
        假设行业数量为N :
        收入分配矩阵 : D (N*N)
        列昂惕夫逆矩阵 : L (N*N)
        最终需求向量 : vector_df (N*1)
        参数:
        name (str): 最终需求项目的名称，必须是 q2 数据框的列名。

        返回:
        pd.Series: 包含每个产业部门的收入诱导系数值的 Series，索引为产业部门名称。
        """
        # D为收入分配系数矩阵
        # L为列昂惕夫逆矩阵
        # vector_df为最终需求向量
        name = '需求向量'
        D = self.income_coef(production_sub_items)
        L = self.L
        vector_df = pd.DataFrame(0, index=self.A.index, columns=[name])
        vector_df.loc[sector, name] = 1
        return D.dot(L).dot(vector_df)

    def PIC_all(self):
        """
        计算不同最终需求项目对各产业部门生产活动的总拉动值。

        返回:
        pd.Series: 包含不同最终需求项目总拉动值的 Series，索引为最终需求项目名称，名称为 '不同最终需求的拉动'。
        """
        # 初始化一个空字典，用于存储不同最终需求项目的总拉动值
        dic = {}
        # 遍历不同的最终需求项目名称列表
        for name in [
            '农村居民消费支出',
            '城镇居民消费支出',
            '居民消费支出',
            '固定资本形成总额',
            '资本形成总额',
            '出口',
            '最终使用合计'
            ]:
            # 调用 PIC_q2 方法计算指定最终需求项目的生产诱发度，
            # 对结果求和并获取其第一个值，将其作为该最终需求项目的总拉动值存入字典
            dic[name] = self.PIC_q2(name).sum().values[0]
        # 将字典转换为 Pandas Series 对象并返回，指定 Series 的名称为 '不同最终需求的拉动'
        return pd.Series(dic, name='不同最终需求的拉动')


    def IIC_all(self, production_sub_items='劳动者报酬'):
        """
        计算不同产业部门在指定生产子项上的收入诱发度总和。

        参数:
        production_sub_items (str): 生产子项名称，默认为 '劳动者报酬'。

        返回:
        pd.Series: 包含不同产业部门收入诱发度总和的 Series，索引为产业部门名称，
                   名称为 '不同最终需求对{production_sub_items}的拉动'。
        """
        # 初始化一个空字典，用于存储不同产业部门的收入诱发度总和
        dic = {}
        # 遍历直接消耗系数矩阵 self.A 的所有行索引，即所有产业部门
        for sector in self.A.index:
            # 调用 IIC_sector 方法计算指定产业部门在指定生产子项上的收入诱发度，
            # 并获取结果 DataFrame 中的第一个元素，将其作为该产业部门的收入诱发度总和存入字典
            dic[sector] = self.IIC_sector(sector, production_sub_items).values[0][0]
        # 将字典转换为 Pandas Series 对象并返回，指定 Series 的名称
        return pd.Series(dic, name=f'不同最终需求对{production_sub_items}的拉动')


class MultiYearInputOutputTable:
    def __init__(self, year_list = [2002, 2007, 2012, 2017, 2020, 2023]):
        self.year_list = year_list
        self.io_data = {}
        for year in year_list:
            data = pd.read_excel(rf'data/中国投入产出表/{year}.xlsx', sheet_name=None, index_col=0)
            io = InputOutputTable(year, '中国', data, '40行业')
            setattr(self, f"Y{year}", io)
            self.io_data[year] = io

    def compare_items(self, q):
        # 使用字典推导式创建一个字典，键为年份，值为对应年份 InputOutputTable 对象的指定属性 q 的值
        # self.io_data 是一个字典，键为年份，值为对应的 InputOutputTable 实例
        # getattr(io, q) 用于获取 io 对象的属性 q 的值
        df = pd.DataFrame({year:getattr(io, q) for year, io in self.io_data.items()})
        # 将包含不同年份指定属性值的字典转换为 DataFrame 并返回
        return df

    def plot_multi_year_items(self, q):
        """
        绘制多年份下指定属性数据的分组条形图。

        参数:
        q (str): 要绘制的数据的属性名，用于从 MultiYearInputOutputTable 实例中获取对应数据。

        返回:
        plotly.graph_objects.Figure: 包含分组条形图的 Plotly 图形对象。
        """
        # 调用 compare_items 方法获取多年份的指定属性数据的 DataFrame
        # 并使用 Plotly Express 的 bar 函数创建分组条形图
        fig = px.bar(self.compare_items(q), barmode='group',
                     # 设置坐标轴和图例标签，横坐标显示行业，纵坐标显示金额单位为亿元，图例显示投入产出年份
                     labels={'index':'行业', 'value':'亿元', 'variable':'投入产出年份'})
        return fig

    def plot_multi_year_items_ratio(self, numerator, denominator):
        """
        绘制多年份下指定分子与分母数据比值的分组条形图。

        参数:
        numerator (str): 分子数据的属性名，用于从 MultiYearInputOutputTable 实例中获取对应数据。
        denominator (str): 分母数据的属性名，用于从 MultiYearInputOutputTable 实例中获取对应数据。

        返回:
        plotly.graph_objects.Figure: 包含分组条形图的 Plotly 图形对象。
        """
        # 获取分子数据，通过 compare_items 方法获取多年份的分子数据 DataFrame
        df_num = self.compare_items(numerator)
        # 获取分母数据，通过 compare_items 方法获取多年份的分母数据 DataFrame
        df_den = self.compare_items(denominator)
        # 计算比值，将分子数据 DataFrame 除以分母数据 DataFrame 得到比值 DataFrame
        df = df_num / df_den

        # 绘制分组条形图，使用 Plotly Express 的 bar 函数创建图表
        fig = px.bar(df, barmode='group',
                     # 设置坐标轴和图例标签，横坐标显示行业及分子分母信息，纵坐标显示比例，图例显示投入产出年份
                     labels={'index':f'行业{numerator}/{denominator}', 'value':'比例', 'variable':'投入产出年份'})
        return fig

# 示例使用
if __name__ == "__main__":
    data2020 = pd.read_excel(r'data/中国投入产出表/2020.xlsx', sheet_name=None, index_col=0)
    io20 = InputOutputTable(2023, '中国', data2020, '40行业')
    
    mio = MultiYearInputOutputTable()
